package org.jeecg.common.es.template;

import org.elasticsearch.client.RestHighLevelClient;
import org.jeecg.common.es.EsIndexUtil;
import org.jeecg.common.es.entity.EsEntity;
import org.jeecg.common.exception.JeecgBootException;
import org.jeecg.common.util.DateUtils;
import org.springframework.data.elasticsearch.core.ElasticsearchRestTemplate;
import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverter;
import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentEntity;
import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;

import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * @author wangzun1.vendor@sensetime.com
 * @Description:
 * @date 2022/8/10
 */
public class CssdElasticsearchRestTemplate extends ElasticsearchRestTemplate {
    public static final String SEPARATOR = "_";
    public static final String PARAMS_POINT = "params.";
    public static final String CTX_SOURCE = "ctx._source.";

    public CssdElasticsearchRestTemplate(RestHighLevelClient client) {
        super(client);
    }

    public CssdElasticsearchRestTemplate(RestHighLevelClient client, ElasticsearchConverter elasticsearchConverter) {
        super(client, elasticsearchConverter);
    }

    @Override
    public <T> T save(T entity) {
        String tenantId = EsIndexUtil.getTenantId();
        if (entity instanceof EsEntity) {
            EsEntity esEntity = (EsEntity)entity;
            Long createTime = esEntity.getSpiderTime();
            String prefixIndexName = this.getPrefixIndexName(esEntity.getClass());
            IndexCoordinates indexCoordinates = IndexCoordinates
                .of(new String[] {prefixIndexName + "_" + tenantId + "_" + DateUtils.formatDate( createTime,"yyyyMM")});
            return super.save(entity, indexCoordinates);
        } else {
            throw new JeecgBootException("传入对象不是[EsEntity]或没有继承[EsEntity]!");
        }
    }

    @Override
    public <T> Iterable<T> save(Iterable<T> entities) {
        String tenantId = EsIndexUtil.getTenantId();
        if (entities.iterator().next() instanceof EsEntity) {
            List<T> saveList = (List)entities;
            Map<String, List<T>> collect = (Map)saveList.stream().collect(Collectors.groupingBy((value) -> {
                return DateUtils.formatDate( ((EsEntity)value).getSpiderTime(),"yyyyMM");
            }));
            List<T> result = new LinkedList();
            collect.forEach((key, value) -> {
                EsEntity esEntity = (EsEntity)value.get(0);
                String prefixIndexName = this.getPrefixIndexName(esEntity.getClass());
                IndexCoordinates indexCoordinates = IndexCoordinates.of(new String[] {
                    prefixIndexName + "_" + tenantId + "_" + DateUtils.formatDate( esEntity.getSpiderTime(),"yyyyMM")});
                Iterable<T> save = super.save(value, indexCoordinates);
                result.addAll((List)save);
            });
            return result;
        } else {
            throw new JeecgBootException( "传入集合元素不是[EsEntity]或没有继承[EsEntity]!");
        }
    }

    public String getPrefixIndexName(Class<? extends EsEntity> clazz) {
        String prefixIndexName = ((ElasticsearchPersistentEntity)Objects
            .requireNonNull(this.getElasticsearchConverter().getMappingContext().getPersistentEntity(clazz)))
                .getIndexCoordinates().getIndexName();
        int lastIndexOf = prefixIndexName.trim().lastIndexOf("_");
        if (lastIndexOf == prefixIndexName.length() - 1) {
            prefixIndexName = prefixIndexName.substring(0, lastIndexOf);
        }

        return prefixIndexName;
    }
}
